package code.limiter;

import java.time.Duration;
import java.util.*;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

public class SlidingWindow {
    private long capacity; // 总大小限制
    private long window; // 大窗口
    private long smallWindow; // 小窗口
    private long smallWindows; // 小窗口数
    private Map<Long, Long> counter; // 窗口计数器
    private final Lock lock = new ReentrantLock();
    private final Object lock2 = new Object();

    public SlidingWindow(long _capacity, long _window, long _smallWindows) {
        capacity = _capacity;
        window = _window;
        smallWindow = _smallWindows;
        smallWindows = _window / _smallWindows;
        counter = new HashMap<>();
    }

    public boolean Allow() {
        lock.lock();
        try {
            // 当前窗口值
            long now = System.nanoTime();
            long currentWindow = now / smallWindow * smallWindow;
            // 起始窗口值
            long startWindow = currentWindow - (smallWindows - 1) * smallWindow;

            // 计算滑动窗口内的总大小，并且移除过期的窗口
            long sumCount = 0;
            for (long key : counter.keySet()) {
                if (key < startWindow) {
                    counter.remove(key);
                    continue;
                }
                sumCount += counter.get(key);
            }
            if (sumCount >= capacity) {
                return false;
            }
            if (counter.containsKey(currentWindow)) {
                long cnt = counter.getOrDefault(currentWindow, Long.valueOf(0)) + 1;
                counter.put(currentWindow, cnt);
            }else {
                counter.put(currentWindow, Long.valueOf(1));
            }
            return true;
        }finally {
            lock.unlock();
        }
    }

    public static void main(String[] args) {
        Duration window = Duration.ofSeconds(1);
        Duration smallWindow = Duration.ofMillis(100);
        SlidingWindow limiter = new SlidingWindow(3, window.toNanos(), smallWindow.toNanos()); // 控制1秒钟3个请求
        for (int i=0; i<10; i++) {
            if (limiter.Allow()) {
                System.out.println(i+": OK");
            }else {
                System.out.println(i+": False");
            }
        }
    }
}
